package com.ajocer.springbootinit.config;

import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.param.partition.CreatePartitionParam;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.annotation.PostConstruct;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import com.google.gson.Gson;

/**
 * Milvus高级配置类
 * 优化向量检索性能
 */
@Configuration
@Slf4j
public class MilvusAdvancedConfig {

    @Autowired
    private MilvusServiceClient milvusClient;

    @Autowired
    private MilvusConfig milvusConfig;

    /**
     * 初始化高级配置
     */
    @PostConstruct
    public void init() {
        try {
            // 创建分区，按景点类型分区，提高查询效率
            createPartitions();
            
            // 优化索引参数
            optimizeIndexParams();
            
            log.info("Milvus高级配置初始化完成");
        } catch (Exception e) {
            log.error("Milvus高级配置初始化失败", e);
        }
    }

    /**
     * 创建分区
     * 按景点类型分区，提高查询效率
     */
    private void createPartitions() {
        String collectionName = milvusConfig.getCollection();
        
        // 常见景点类型
        String[] spotTypes = {
                "historical", // 历史文化
                "natural",    // 自然风光
                "entertainment", // 娱乐
                "food",       // 美食
                "shopping"    // 购物
        };
        
        // 为每种类型创建分区
        for (String type : spotTypes) {
            try {
                CreatePartitionParam createPartitionParam = CreatePartitionParam.newBuilder()
                        .withCollectionName(collectionName)
                        .withPartitionName("partition_" + type)
                        .build();
                
                // 检查分区是否存在
                R<Boolean> response = milvusClient.hasPartition(
                        io.milvus.param.partition.HasPartitionParam.newBuilder()
                                .withCollectionName(collectionName)
                                .withPartitionName("partition_" + type)
                                .build()
                );
                if (!response.getData()) {
                    milvusClient.createPartition(createPartitionParam);
                    log.info("创建分区成功: partition_{}", type);
                } else {
                    log.info("分区已存在: partition_{}", type);
                }
            } catch (Exception e) {
                log.warn("创建分区失败: partition_{}", type, e);
            }
        }
    }

    /**
     * 优化索引参数
     */
    private void optimizeIndexParams() {
        String collectionName = milvusConfig.getCollection();
        String fieldName = "embedding";
        
        try {
            // 为HNSW索引设置优化参数
            Map<String, Object> extraParams = new HashMap<String, Object>();
            extraParams.put("M", 16);          // 每个节点的最大连接数
            extraParams.put("efConstruction", 200); // 构建索引时的搜索宽度
            
            CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withFieldName(fieldName)
                    .withIndexType(milvusConfig.getIndexType())
                    .withMetricType(milvusConfig.getMetricType())
                    .withExtraParam(new Gson().toJson(extraParams))
                    .withSyncMode(Boolean.TRUE)
                    .build();
            
            milvusClient.createIndex(createIndexParam);
            log.info("索引参数优化完成");
            
            // 重新加载集合，使配置生效
            milvusClient.loadCollection(
                    LoadCollectionParam.newBuilder()
                            .withCollectionName(collectionName)
                            .build()
            );
        } catch (Exception e) {
            log.warn("索引参数优化失败", e);
        }
    }

    /**
     * 创建优化的搜索参数构建器
     */
    @Bean
    public SearchParam.Builder searchParamBuilder() {
        // 优化搜索参数
        Map<String, Object> searchParams = new HashMap<String, Object>();
        searchParams.put("ef", 200);     // 搜索时的候选集大小
        searchParams.put("nprobe", 16);  // 搜索时探测的聚类数量
        
        return SearchParam.newBuilder()
                .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                .withParams(new Gson().toJson(searchParams))
                .withMetricType(milvusConfig.getMetricType());
    }
} 